#include <ros/ros.h>
#include <geometry_msgs/Point.h>
#include <yolo_detect/YoloResult.h>
#include <traj_utils/DetectPose.h>
#include <eigen3/Eigen/Core>
#include <eigen3/Eigen/Geometry>
#include <vector>
#include "tracker/sord3d.hpp"

class ObjectTracker3D
{
public:
  explicit ObjectTracker3D(ros::NodeHandle &nh) : nh_(nh) {

    nh_.param("camera/color/fx", fx_, -1.0);
    nh_.param("camera/color/fy", fy_, -1.0);
    nh_.param("camera/color/cx", cx_, -1.0);
    nh_.param("camera/color/cy", cy_, -1.0);

    std::cout << "[YOLO_DETECT] fx: " << fx_ << std::endl;
    std::cout << "[YOLO_DETECT] fy: " << fy_ << std::endl;
    std::cout << "[YOLO_DETECT] cx: " << cx_ << std::endl;
    std::cout << "[YOLO_DETECT] cy: " << cy_ << std::endl;

    int min_hits;
    double max_age, dist_thresh;
    nh_.param("tracker/Sord3d/max_age", max_age, -1.0);
    nh_.param("tracker/Sord3d/min_hits", min_hits, -1);
    nh_.param("tracker/Sord3d/dist_thresh", dist_thresh, -1.0);

    std::cout << "[SORT3D] max_age    : " << max_age << std::endl;
    std::cout << "[SORT3D] min_hits   : " << min_hits << std::endl;
    std::cout << "[SORT3D] dist_thresh: " << dist_thresh << std::endl;
    sort3d_.init(max_age, min_hits, dist_thresh);

    detect_sub_ = nh_.subscribe(
        "/perception/yolo_result", 1, &ObjectTracker3D::detectCallback, this);
    track_pub_ =
        nh_.advertise<traj_utils::DetectPose>("/perception/track_result", 1);

    last_time_ = ros::Time::now();
  }

  void detectCallback(const yolo_detect::YoloResult::ConstPtr &msg) {
    Eigen::Vector3d p_cam = {msg->cam_pose.pose.position.x,
                             msg->cam_pose.pose.position.y,
                             msg->cam_pose.pose.position.z};
    Eigen::Quaterniond q_cam(msg->cam_pose.pose.orientation.w,
                             msg->cam_pose.pose.orientation.x,
                             msg->cam_pose.pose.orientation.y,
                             msg->cam_pose.pose.orientation.z);

    std::vector<Eigen::Vector3d> points(msg->poses.size());
    for (size_t i = 0; i < msg->poses.size(); i++) {
      points[i] = p_cam + q_cam * pixelToCamera(msg->poses[i].x,
                                                msg->poses[i].y,
                                                msg->poses[i].z);
    }

    // update the tracker and get the tracking result
    double dt = (msg->header.stamp - last_time_).toSec();
    last_time_ = msg->header.stamp;
    std::vector<Eigen::Vector4d> id_poses = sort3d_.update(points, dt);

    traj_utils::DetectPose track_msg;
    track_msg.header.stamp = msg->header.stamp;
    for (auto &id_pose : id_poses) {
      track_msg.ids.push_back(id_pose(0));
      geometry_msgs::Point point;
      point.x = id_pose(1);
      point.y = id_pose(2);
      point.z = id_pose(3);
      track_msg.poses.push_back(point);
    }

    track_pub_.publish(track_msg);
  }

private:
  Sort3D sort3d_;
  ros::NodeHandle nh_;
  ros::Subscriber detect_sub_;
  ros::Publisher track_pub_;

  ros::Time last_time_{ros::Time(0)};

  // the camera intrinsics
  double fx_{610.8905};
  double fy_{610.8898};
  double cx_{322.7573};
  double cy_{250.6328};

  Eigen::Vector3d pixelToCamera(double x, double y, double depth) {
    double X = (x - cx_) * depth / fx_;
    double Y = (y - cy_) * depth / fy_;
    double Z = depth;
    return {X, Y, Z};
  }
};

int main(int argc, char **argv) {
  ros::init(argc, argv, "object_tracker_3d");
  ros::NodeHandle nh("~");

  ObjectTracker3D tracker(nh);

  ros::spin();
  return 0;
}
